package com.zz.common

import scala.collection.immutable
import scala.concurrent.ExecutionContext

import akka.NotUsed
import akka.actor.{ActorSystem, Scheduler}
import akka.stream.scaladsl.{Broadcast, Flow, GraphDSL, Sink, Source, ZipN}
import akka.stream.{FlowShape, Graph, Materializer}
import cats.effect.IO
import com.typesafe.scalalogging.Logger
import com.zz.common.service.IdService.{Zid, Zidservice}
import com.zz.dsp.idl.DspEvent

/**
  * @author huanggh
  *         19-8-16 上午11:56
  */
object BroadcastFlow {
  def apply(zidService: Zidservice[IO])(implicit ctx: GlobalContext[IO], system: ActorSystem,
              mat: Materializer,
              ec: ExecutionContext): Flow[PassThrough[Seq[DspEvent]], BatchOffsets, NotUsed] = {
    implicit val scheduler: Scheduler = system.scheduler
    new BroadcastFlow[String, Array[Byte], DspEvent] {
      val idBundleFlow: Flow[DspEvent, Zid, NotUsed] = IdBundleFlow(zidService).flow()
    }
  }.flow()
}


trait BroadcastFlow[@specialized K, @specialized V, @specialized E] {

  def idBundleFlow: Flow[E, _, NotUsed]

  protected val logger: Logger = Logger(getClass)

  type P = BatchOffsets

  def flow()(implicit system: ActorSystem, mat: Materializer, ec: ExecutionContext): Flow[PassThrough[Seq[E]], P, NotUsed] = {
    val g: Graph[FlowShape[PassThrough[Seq[E]], immutable.Seq[BatchOffsets]], NotUsed] = GraphDSL.create() { implicit builder =>
      import akka.stream.scaladsl.GraphDSL.Implicits._

      val bc = builder.add(Broadcast[PassThrough[Seq[E]]](1))

      val zip = builder.add(ZipN[P](1))

      def f(flow: Flow[E, _, _]): Flow[PassThrough[Seq[E]], P, NotUsed] = {
        Flow[PassThrough[Seq[E]]]
          .mapAsync(1) { x =>
            Source(x.data.toList).via(flow).runWith(Sink.ignore).map { _ =>
              x.offsets
            }
          }
      }

      val f1: Flow[PassThrough[Seq[E]], P, NotUsed] = f(idBundleFlow)

      bc ~> f1 ~> zip

      FlowShape.apply(bc.in, zip.out)
    }
    val cc: Flow[Seq[P], P, NotUsed] =
      Flow[Seq[P]].mapConcat(_.headOption.toList)
    Flow
      .fromGraph(g)
      .via(cc)
  }
}
